from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from typing import List, Dict
import asyncio

app = FastAPI(title="Vue3 + FastAPI WebSocket")

# 解决跨域问题
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:3000"],  # Vue 前端地址
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 存储所有活跃的 WebSocket 连接
active_connections: List[WebSocket] = []
# 存储客户端ID与连接的映射（可选，用于定向推送）
client_connections: Dict[str, WebSocket] = {}


# 广播消息给所有连接的客户端
async def broadcast(message: str):
    for connection in active_connections:
        await connection.send_text(message)


# 定向推送消息给指定客户端
async def send_to_client(client_id: str, message: str):
    if client_id in client_connections:
        await client_connections[client_id].send_text(message)


@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    # 接受客户端连接
    await websocket.accept()
    # 存储连接
    active_connections.append(websocket)
    client_connections[client_id] = websocket
    print(f"客户端 {client_id} 已连接，当前连接数：{len(active_connections)}")

    try:
        while True:
            # 接收客户端发送的消息
            data = await websocket.receive_text()
            print(f"收到客户端 {client_id} 的消息：{data}")

            # 1. 回复当前客户端（单独回应）
            await websocket.send_text(f"服务器已收到你的消息：{data}")

            # 2. 广播给所有客户端（可选）
            # await broadcast(f"客户端 {client_id} 说：{data}")

            # 3. 模拟服务器主动推送（5秒后推送一条消息）
            await asyncio.sleep(5)
            await websocket.send_text(f"服务器主动推送：你5秒前发送了「{data}」")

    except WebSocketDisconnect:
        # 客户端断开连时清理连接
        active_connections.remove(websocket)
        client_connections.pop(client_id, None)
        print(f"客户端 {client_id} 已断开，当前连接数：{len(active_connections)}")


# 提供一个HTTP接口，用于外部触发服务器推送（可选）
@app.post("/push/{client_id}")
async def push_to_client(client_id: str, message: str):
    for client_id in client_connections:
        await send_to_client(client_id, f"管理员推送：{message}")
        return {"status": "success", "message": f"已向 {client_id} 推送消息"}
    return {"status": "error", "message": f"客户端 {client_id} 未连接"}


import uvicorn

if __name__ == '__main__':
    uvicorn.run("twebsocket:app", host="127.0.0.1", port=8000, reload=True)